from pathlib import Path
import numpy as np
import torch
import yaml
import math
import os
import time
from mynn.datasets import build_dataset, build_dataloader
from mynn.models import build_model
from mynn.utils import tensor2img, load_checkpint, get_root_logger
from PIL import Image
from torch.nn import functional as F

from mynn.utils.logger import log_print

# Read yaml file.
YAML_PATH = 'options/TOFlow/toflow.yaml'
with open(YAML_PATH, 'r', encoding='utf-8') as f:
    opt = yaml.load(f, Loader=yaml.SafeLoader)

# Set logger.
log_dir = Path('experiments') / opt['exp_name']
os.makedirs(log_dir, exist_ok=True)
log_file = log_dir / f"{opt['test']['log_file']}"
_logger = get_root_logger(log_file=log_file)

# Choose CUDA or CPU.
device = "cuda" if opt['test']['cuda'] else "cpu"

# Set GPU list.
gpu_list = ",".join([str(v) for v in opt['test']['gpu_list']])
os.environ['CUDA_VISIBLE_DEVICES'] = gpu_list
log_print(f'device:{device}')

# Get dataset and dataloader.
test_dataset = build_dataset(dataset_opt=opt['test']['dataset'], phase='test')

test_dataloader = build_dataloader(dataset=test_dataset, opt=opt, phase='test')
# Build model.
model = build_model(opt)
model.to(device)
# Load checkpoint.
model, current_iter = load_checkpint(opt=opt, model=model)

log_print('Testing start.')

model.eval()
total_step = len(test_dataloader)
save_root = Path('./experiments') / \
    opt['exp_name'] / 'results' / opt['test']['dataset']['name']
log_print(f'Save path:{save_root}')
for step, data in enumerate(test_dataloader):
    # Unpack data.
    gt = data['gt'].to(device)
    lqs = data['lqs'].to(device)
    key = data['key'][0]

    # Interpolate to lqs for making its shape as same as gt.
    b, n, c, h, w = lqs.shape
    lq_list = []
    for i in range(n):
        lq = F.interpolate(lqs[:, i, :, :, :], scale_factor=4)
        lq_list.append(lq)
    lqs = torch.stack(lq_list, dim=1)


    # # Crop to one eighth.
    # b, n, c, h, w = lqs.shape
    # lqs = lqs[:, :, :, :h // 2, :w // 4]
    # b, c, h, w = gt.shape
    # gt = gt[:, :, :h // 2, :w // 4]

    # Forward propagation.
    sr = model(lqs)

    # Save result.
    clip_name, image_name = key.split('/')
    save_dir = save_root / clip_name
    os.makedirs(save_dir, exist_ok=True)
    save_path = save_dir / f'{image_name}.png'

    sr_img = tensor2img(sr)
    sr_img = Image.fromarray(sr_img)
    sr_img.save(save_path)
    print(f'{step+1}/{total_step}: {key}')

log_print('Test complete.')
